from collections import Counter
from tool import sys_tool
from loaders import data_iflytek
from loaders import data as data_xiecheng_jiudian

from models import xlnet_sweet_add
from models import A3_resnet_sweet_add
from models import resnet_sweet_add
from models import resnet

from configs import c5_xlnet_f1
from configs import c5A_xlnet_f1
from configs import c5B_xlnet_f1

from learners import learner1_mlp
from learners import learner2_resnet
from learners import learner3A_sweet_add
from learners import learner3B_sweet_add
from learners import learner4_xlnet
from learners import learner5_xlnet_f1

# 数据加载器
loaders = {
    'data_xiecheng_jiudian': data_xiecheng_jiudian.get_data,  # 携程网_酒店_评价分类
    'data_iflytek': data_iflytek.get_data,  # APP 类别预测数据
}

# 模型构造器
models = {
    'xlnet': xlnet_sweet_add.create_model,
    'A3_resnet_sweet_add': A3_resnet_sweet_add.create_model,
    'resnet_sweet_add': resnet_sweet_add.create_model,
    'resnet': resnet.create_model,
}

# 学习器
learners = {
    'learner1_mlp': learner1_mlp.Learner,
    'learner2_resnet': learner2_resnet.Learner,
    'learner3A_sweet_add': learner3A_sweet_add.Learner,
    'learner3B_sweet_add': learner3B_sweet_add.Learner,
    'learner4_xlnet': learner4_xlnet.Learner,
    'learner5_xlnet_f1': learner5_xlnet_f1.Learner,
}

# 配置
configs = {
    'c5': c5_xlnet_f1.Config,
    'c5A': c5A_xlnet_f1.Config,
    'c5B': c5B_xlnet_f1.Config,
}

# 检查不同模块注册 不要重名
c = Counter()
for kv in [loaders, models, learners, configs]:
    c.update(list(kv))
same_names = [k for k, v in c.items() if v > 1]
if same_names:
    raise sys_tool.InfoError(f'重名={same_names}')


# 根据名字获取注册对象
def get(name):
    for kv in [loaders, models, learners, configs]:
        if name in kv:
            return kv[name]
    raise sys_tool.InfoError(f'未发现注册的数据加载器、模型构造器、配置 name={name}')
